#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
回测系统
用于回测和评估低位龙头股分析策略的有效性
"""

import os
import json
import logging
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
from pathlib import Path

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

class BackTester:
    """回测系统，用于验证低位龙头股分析策略的有效性"""
    
    def __init__(self, data_path='data'):
        """
        初始化回测系统
        
        Args:
            data_path: 数据目录路径
        """
        self.data_path = data_path
        self.price_data_path = os.path.join(data_path, 'price_history')
        self.results_path = os.path.join(data_path, 'results', 'backtesting')
        
        # 创建结果目录
        os.makedirs(self.results_path, exist_ok=True)
        
        # 回测结果
        self.performance_metrics = {}
        
    def load_stock_prices(self, stock_code):
        """
        加载股票历史价格数据
        
        Args:
            stock_code: 股票代码
            
        Returns:
            pd.DataFrame: 股票价格数据
        """
        try:
            # 构建文件路径
            file_path = os.path.join(self.price_data_path, f"{stock_code}.json")
            
            if not os.path.exists(file_path):
                logger.warning(f"未找到股票 {stock_code} 的历史价格数据")
                return pd.DataFrame()
                
            # 读取JSON文件
            with open(file_path, 'r', encoding='utf-8') as f:
                price_data = json.load(f)
                
            # 检查数据结构
            if 'daily_prices' not in price_data or not price_data['daily_prices']:
                logger.warning(f"股票 {stock_code} 的历史价格数据格式错误或为空")
                return pd.DataFrame()
                
            # 转换为DataFrame
            df = pd.DataFrame(price_data['daily_prices'])
            
            # 确保必要的列存在
            required_columns = ['date', 'open', 'high', 'low', 'close', 'vol']
            for col in required_columns:
                if col not in df.columns:
                    logger.warning(f"股票 {stock_code} 的历史价格数据缺少必要列: {col}")
                    return pd.DataFrame()
                    
            # 转换日期列并排序
            df['date'] = pd.to_datetime(df['date'])
            df = df.sort_values('date').reset_index(drop=True)
            
            # 确保数值列为浮点型
            numeric_columns = ['open', 'high', 'low', 'close', 'vol']
            for col in numeric_columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')
                
            return df
            
        except Exception as e:
            logger.error(f"加载股票 {stock_code} 的历史价格数据时出错: {e}")
            return pd.DataFrame()
            
    def load_analysis_result(self, result_file):
        """
        加载低位龙头股分析结果
        
        Args:
            result_file: 结果文件路径
            
        Returns:
            list: 分析结果列表
        """
        try:
            with open(result_file, 'r', encoding='utf-8') as f:
                result = json.load(f)
                
            return result
            
        except Exception as e:
            logger.error(f"加载分析结果失败: {e}")
            return []
            
    def calculate_returns(self, stock_prices, start_date, periods=[1, 3, 6, 12], unit='month'):
        """
        计算不同时间段的收益率
        
        Args:
            stock_prices: 股票价格DataFrame
            start_date: 开始日期
            periods: 时间段列表，单位由unit参数指定
            unit: 时间单位，可选 'day', 'week', 'month'
            
        Returns:
            dict: 各时间段的收益率
        """
        try:
            # 转换日期格式
            if isinstance(start_date, str):
                start_date = pd.to_datetime(start_date)
                
            # 确保数据已按日期排序
            stock_prices = stock_prices.sort_values('date')
            
            # 找到开始日期对应的价格
            start_prices = stock_prices[stock_prices['date'] >= start_date]
            if start_prices.empty:
                logger.warning(f"未找到 {start_date} 之后的价格数据")
                return {period: 0 for period in periods}
                
            start_price = start_prices.iloc[0]['close']
            start_idx = start_prices.index[0]
            
            # 计算各时间段的收益率
            returns = {}
            
            for period in periods:
                # 计算结束日期
                if unit == 'day':
                    end_date = start_date + timedelta(days=period)
                elif unit == 'week':
                    end_date = start_date + timedelta(weeks=period)
                elif unit == 'month':
                    # 简化处理：一个月按30天计算
                    end_date = start_date + timedelta(days=period*30)
                else:
                    raise ValueError(f"不支持的时间单位: {unit}")
                    
                # 找到最接近结束日期的价格数据
                end_prices = stock_prices[stock_prices['date'] >= end_date]
                
                if end_prices.empty:
                    # 如果没有对应的结束价格，使用最新价格
                    end_price = stock_prices.iloc[-1]['close']
                else:
                    end_price = end_prices.iloc[0]['close']
                    
                # 计算收益率
                returns[period] = (end_price / start_price - 1) * 100
                
            return returns
            
        except Exception as e:
            logger.error(f"计算收益率失败: {e}")
            return {period: 0 for period in periods}
            
    def backtest_single_result(self, result_file, periods=[1, 3, 6, 12], unit='month'):
        """
        回测单个分析结果
        
        Args:
            result_file: 结果文件路径
            periods: 时间段列表
            unit: 时间单位
            
        Returns:
            dict: 回测结果
        """
        try:
            # 提取日期信息
            file_name = os.path.basename(result_file)
            date_str = file_name.split('_')[-1].split('.')[0]
            result_date = datetime.strptime(date_str, "%Y%m%d")
            
            # 加载分析结果
            results = self.load_analysis_result(result_file)
            if not results:
                logger.warning(f"分析结果为空: {result_file}")
                return None
                
            # 回测结果
            backtest_result = {
                'date': result_date.strftime("%Y-%m-%d"),
                'stocks': [],
                'average_returns': {period: 0 for period in periods},
                'win_rate': {period: 0 for period in periods},
                'benchmark_returns': {period: 0 for period in periods}
            }
            
            # 加载指数数据作为基准
            benchmark_code = '000001'  # 上证指数
            benchmark_prices = self.load_stock_prices(benchmark_code)
            if not benchmark_prices.empty:
                benchmark_returns = self.calculate_returns(benchmark_prices, result_date, periods, unit)
                backtest_result['benchmark_returns'] = benchmark_returns
            
            # 计算每只股票的收益率
            valid_periods_count = {period: 0 for period in periods}
            win_count = {period: 0 for period in periods}
            
            for stock in results:
                stock_code = stock['code']
                stock_name = stock['name']
                
                # 加载股票价格数据
                stock_prices = self.load_stock_prices(stock_code)
                if stock_prices.empty:
                    logger.warning(f"未找到股票 {stock_code} 的价格数据")
                    continue
                    
                # 计算各时间段的收益率
                returns = self.calculate_returns(stock_prices, result_date, periods, unit)
                
                # 统计胜率
                for period, ret in returns.items():
                    if ret > 0:
                        win_count[period] += 1
                    valid_periods_count[period] += 1
                    backtest_result['average_returns'][period] += ret
                
                # 记录股票收益率
                stock_result = {
                    'code': stock_code,
                    'name': stock_name,
                    'returns': returns,
                    'score': stock.get('total_score', 0)
                }
                backtest_result['stocks'].append(stock_result)
            
            # 计算平均收益率和胜率
            for period in periods:
                if valid_periods_count[period] > 0:
                    backtest_result['average_returns'][period] /= valid_periods_count[period]
                    backtest_result['win_rate'][period] = (win_count[period] / valid_periods_count[period]) * 100
            
            return backtest_result
            
        except Exception as e:
            logger.error(f"回测 {result_file} 失败: {e}")
            return None
            
    def backtest_all_results(self, result_dir=None):
        """
        回测所有分析结果
        
        Args:
            result_dir: 结果目录路径，默认为data/results/low_position_leaders
            
        Returns:
            list: 回测结果列表
        """
        if result_dir is None:
            result_dir = os.path.join(self.data_path, 'results', 'low_position_leaders')
            
        try:
            # 检查目录是否存在
            if not os.path.exists(result_dir):
                logger.warning(f"结果目录不存在: {result_dir}")
                return []
                
            # 找到所有结果文件
            result_files = []
            for file in os.listdir(result_dir):
                if file.startswith('low_position_leaders_') and file.endswith('.json'):
                    result_files.append(os.path.join(result_dir, file))
                    
            if not result_files:
                logger.warning(f"未找到分析结果文件")
                return []
                
            # 按日期排序
            result_files.sort()
            
            # 回测结果
            backtest_results = []
            
            # 回测每个结果
            for result_file in result_files:
                logger.info(f"回测分析结果: {result_file}")
                result = self.backtest_single_result(result_file)
                if result:
                    backtest_results.append(result)
                    
            return backtest_results
            
        except Exception as e:
            logger.error(f"回测所有结果失败: {e}")
            return []
            
    def calculate_performance_metrics(self, backtest_results):
        """
        计算回测性能指标
        
        Args:
            backtest_results: 回测结果列表
            
        Returns:
            dict: 性能指标
        """
        try:
            if not backtest_results:
                logger.warning("回测结果为空，无法计算性能指标")
                return {}
                
            # 提取所有时间段
            periods = list(backtest_results[0]['average_returns'].keys())
            
            # 计算各时间段的平均收益率、最大收益率、最小收益率
            performance = {}
            
            for period in periods:
                returns = [result['average_returns'][period] for result in backtest_results]
                benchmark_returns = [result['benchmark_returns'][period] for result in backtest_results]
                win_rates = [result['win_rate'][period] for result in backtest_results]
                
                # 超额收益
                excess_returns = [r - b for r, b in zip(returns, benchmark_returns)]
                
                performance[period] = {
                    'avg_return': np.mean(returns),
                    'max_return': np.max(returns),
                    'min_return': np.min(returns),
                    'std_return': np.std(returns),
                    'avg_benchmark_return': np.mean(benchmark_returns),
                    'avg_excess_return': np.mean(excess_returns),
                    'avg_win_rate': np.mean(win_rates)
                }
                
            # 计算总体性能指标
            dates = [result['date'] for result in backtest_results]
            
            # 记录性能指标
            self.performance_metrics = {
                'periods': performance,
                'dates': dates,
                'total_tests': len(backtest_results)
            }
            
            return self.performance_metrics
            
        except Exception as e:
            logger.error(f"计算性能指标失败: {e}")
            return {}
            
    def visualize_performance(self, output_path=None):
        """
        可视化回测性能
        
        Args:
            output_path: 输出路径，默认为results_path/performance.png
        """
        try:
            if not self.performance_metrics:
                logger.warning("性能指标为空，无法可视化")
                return
                
            if output_path is None:
                output_path = os.path.join(self.results_path, 'performance.png')
                
            # 设置中文字体
            plt.rcParams['font.sans-serif'] = ['SimHei']  # 用来正常显示中文标签
            plt.rcParams['axes.unicode_minus'] = False  # 用来正常显示负号
            
            # 创建图表
            fig, axes = plt.subplots(2, 2, figsize=(16, 12))
            
            # 提取数据
            periods = list(self.performance_metrics['periods'].keys())
            periods.sort()
            
            # 图1：各时间段平均收益率vs基准收益率
            avg_returns = [self.performance_metrics['periods'][p]['avg_return'] for p in periods]
            benchmark_returns = [self.performance_metrics['periods'][p]['avg_benchmark_return'] for p in periods]
            
            ax1 = axes[0, 0]
            width = 0.35
            x = range(len(periods))
            ax1.bar([i - width/2 for i in x], avg_returns, width, label='策略收益')
            ax1.bar([i + width/2 for i in x], benchmark_returns, width, label='基准收益')
            ax1.set_xticks(x)
            ax1.set_xticklabels([f"{p}个月" for p in periods])
            ax1.set_ylabel('收益率(%)')
            ax1.set_title('各时间段平均收益率对比')
            ax1.legend()
            
            # 图2：各时间段平均超额收益率
            excess_returns = [self.performance_metrics['periods'][p]['avg_excess_return'] for p in periods]
            
            ax2 = axes[0, 1]
            ax2.bar(x, excess_returns, color='green')
            ax2.set_xticks(x)
            ax2.set_xticklabels([f"{p}个月" for p in periods])
            ax2.set_ylabel('超额收益率(%)')
            ax2.set_title('各时间段平均超额收益率')
            
            # 图3：各时间段胜率
            win_rates = [self.performance_metrics['periods'][p]['avg_win_rate'] for p in periods]
            
            ax3 = axes[1, 0]
            ax3.bar(x, win_rates, color='orange')
            ax3.set_xticks(x)
            ax3.set_xticklabels([f"{p}个月" for p in periods])
            ax3.set_ylabel('胜率(%)')
            ax3.set_title('各时间段平均胜率')
            
            # 图4：收益率分布情况
            typical_period = periods[len(periods)//2]  # 选择中间的时间段
            returns = [self.performance_metrics['periods'][typical_period]['avg_return'] for _ in self.performance_metrics['dates']]
            
            ax4 = axes[1, 1]
            ax4.hist(returns, bins=10, alpha=0.7, color='purple')
            ax4.set_xlabel('收益率(%)')
            ax4.set_ylabel('频次')
            ax4.set_title(f'{typical_period}个月收益率分布')
            
            # 调整布局
            plt.tight_layout()
            
            # 保存图表
            plt.savefig(output_path, dpi=300, bbox_inches='tight')
            plt.close()
            
            logger.info(f"性能可视化已保存至: {output_path}")
            
        except Exception as e:
            logger.error(f"可视化性能指标失败: {e}")
            
    def generate_report(self, output_path=None):
        """
        生成回测报告
        
        Args:
            output_path: 输出路径，默认为results_path/report.json
            
        Returns:
            dict: 回测报告
        """
        try:
            if not self.performance_metrics:
                logger.warning("性能指标为空，无法生成报告")
                return {}
                
            if output_path is None:
                output_path = os.path.join(self.results_path, 'report.json')
                
            # 构建报告
            report = {
                'timestamp': datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
                'performance_metrics': self.performance_metrics,
                'summary': {
                    'total_tests': self.performance_metrics['total_tests'],
                    'date_range': f"{self.performance_metrics['dates'][0]} 至 {self.performance_metrics['dates'][-1]}"
                }
            }
            
            # 添加总结信息
            periods = list(self.performance_metrics['periods'].keys())
            periods.sort()
            
            overall_avg_return = np.mean([self.performance_metrics['periods'][p]['avg_return'] for p in periods])
            overall_excess_return = np.mean([self.performance_metrics['periods'][p]['avg_excess_return'] for p in periods])
            overall_win_rate = np.mean([self.performance_metrics['periods'][p]['avg_win_rate'] for p in periods])
            
            report['summary']['overall_avg_return'] = overall_avg_return
            report['summary']['overall_excess_return'] = overall_excess_return
            report['summary']['overall_win_rate'] = overall_win_rate
            
            # 保存报告
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(report, f, ensure_ascii=False, indent=2)
                
            logger.info(f"回测报告已保存至: {output_path}")
            
            return report
            
        except Exception as e:
            logger.error(f"生成回测报告失败: {e}")
            return {}
            
    def run_backtest(self, result_dir=None):
        """
        运行完整回测流程
        
        Args:
            result_dir: 结果目录路径
            
        Returns:
            dict: 回测报告
        """
        try:
            logger.info("开始回测低位龙头股策略")
            
            # 回测所有结果
            backtest_results = self.backtest_all_results(result_dir)
            
            if not backtest_results:
                logger.warning("回测结果为空")
                return {}
                
            # 计算性能指标
            performance_metrics = self.calculate_performance_metrics(backtest_results)
            
            # 可视化性能
            self.visualize_performance()
            
            # 生成报告
            report = self.generate_report()
            
            logger.info("回测完成")
            return report
            
        except Exception as e:
            logger.error(f"回测失败: {e}")
            return {}


# 用法示例
if __name__ == "__main__":
    backtester = BackTester()
    report = backtester.run_backtest()
    print(f"回测完成，平均收益率: {report.get('summary', {}).get('overall_avg_return', 0):.2f}%") 